import torch
import torch.nn.functional as F
from core.utils.utils import bilinear_sampler

# 组合几何编码体积
class Combined_Geo_Encoding_Volume:
    def __init__(self, init_fmap1, init_fmap2, geo_volume, num_levels=2, radius=4):
        self.num_levels = num_levels # 相关性金字塔层数
        self.radius = radius # 相关性金字塔宽度
        self.geo_volume_pyramid = [] # 几何编码体积金字塔List
        self.init_corr_pyramid = [] # 全对相关性体积金字塔List

        # all pairs correlation
        init_corr = Combined_Geo_Encoding_Volume.corr(init_fmap1, init_fmap2) # 1/4原尺寸全对相关性体积，输出(B,H,W,1,W)

        b, h, w, _, w2 = init_corr.shape
        b, c, d, h, w = geo_volume.shape
        geo_volume = geo_volume.permute(0, 3, 4, 1, 2).reshape(b*h*w, c, 1, d) # 输入1/4原尺寸/最大视差几何编码体积(B,组数,D,H,W)，转换维度为(B,H,W,组数,D)，改变形状为(B*H*W,组数,1,D)

        init_corr = init_corr.reshape(b*h*w, 1, 1, w2) # 改变1/4原尺寸全对相关性体积形状为(B*H*W,1,1,W)
        self.geo_volume_pyramid.append(geo_volume) # 几何编码体积金字塔List，加入几何编码体积
        self.init_corr_pyramid.append(init_corr) # 全对相关性体积金字塔List，加入全对相关性体积
        for i in range(self.num_levels-1): # i=0，range(1)
            geo_volume = F.avg_pool2d(geo_volume, [1,2], stride=[1,2]) # 对几何编码体积做平均池化，输出1/4原尺寸(B*H*W,组数,1,D/2)
            self.geo_volume_pyramid.append(geo_volume) # 几何编码体积金字塔List，加入视差维度平均池化1/2的几何编码体积

        for i in range(self.num_levels-1):
            init_corr = F.avg_pool2d(init_corr, [1,2], stride=[1,2]) # 对全对相关性体积做平均池化，输出1/4原尺寸(B*H*W,1,1,W/2)
            self.init_corr_pyramid.append(init_corr) # 全对相关性体积金字塔List，加入宽度维度平均池化1/2的全对相关性体积

    def __call__(self, disp, coords):
        r = self.radius # 相关性金字塔宽度=4
        b, _, h, w = disp.shape
        out_pyramid = []
        for i in range(self.num_levels): # 相关性金字塔层数=2
            geo_volume = self.geo_volume_pyramid[i] # 几何编码体积，(B*H*W,组数,1,D)，(B*H*W,组数,1,D/2)
            dx = torch.linspace(-r, r, 2*r+1) # 生成均匀分布1维张量，输出(2*r+1)
            dx = dx.view(1, 1, 2*r+1, 1).to(disp.device) # 调整张量形状为(1,1,2*r+1,1)
            x0 = dx + disp.reshape(b*h*w, 1, 1, 1) / 2**i # (1,1,2*r+1,1)+(B*H*W,1,1,1)/2^i，形状不同张量相加，使用广播规则先扩展到相同维度再相加，输出(B*H*W,1,2*r+1,1)
            y0 = torch.zeros_like(x0) # 生成x0相同形状的零张量

            disp_lvl = torch.cat([x0,y0], dim=-1) # 在最后一个维度连接张量，输出(B*H*W,1,2*r+1,2)
            geo_volume = bilinear_sampler(geo_volume, disp_lvl) # 双线性插值，输出(B*H*W,组数,1,9)
            geo_volume = geo_volume.view(b, h, w, -1) # 几何编码体积，调整张量形状为(B,H,W,72)，-1表示根据张量总大小自动计算最后一个维度的大小=8(组数)*9

            init_corr = self.init_corr_pyramid[i] # 全对相关性体积，(B*H*W,1,1,W)，(B*H*W,1,1,W/2)
            init_x0 = coords.reshape(b*h*w, 1, 1, 1)/2**i - disp.reshape(b*h*w, 1, 1, 1) / 2**i + dx # (B*H*W,1,1,1)/2^i-(B*H*W,1,1,1)/2^i+(1,1,2*r+1,1)，形状不同张量相加，使用广播规则先扩展到相同维度再相加，输出(B*H*W,1,2*r+1,1)

            init_coords_lvl = torch.cat([init_x0,y0], dim=-1) # 在最后一个维度连接张量，输出(B*H*W,1,2*r+1,2)
            init_corr = bilinear_sampler(init_corr, init_coords_lvl) # 双线性插值，输出(B*H*W,1,1,9)
            init_corr = init_corr.view(b, h, w, -1) # 全对相关性体积，调整张量形状为(B,H,W,9)，-1表示根据张量总大小自动计算最后一个维度的大小=9

            out_pyramid.append(geo_volume) # 输出金字塔List，加入几何编码体积
            out_pyramid.append(init_corr) # 输出金字塔List，加入全对相关性
        out = torch.cat(out_pyramid, dim=-1) # 在最后一个维度连接张量，输出(B,H,W,162)，162=几何72+全对9+几何72+全对9
        return out.permute(0, 3, 1, 2).contiguous().float() # 组合几何编码体积(1/4原尺寸)，转换维度为(B,162,H,W)，3d体积

    # 计算全对相关性
    @staticmethod
    def corr(fmap1, fmap2):
        B, D, H, W1 = fmap1.shape # 左右特征图(1/4原尺寸)
        _, _, _, W2 = fmap2.shape
        fmap1 = fmap1.view(B, D, H, W1)
        fmap2 = fmap2.view(B, D, H, W2)
        corr = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2) # 爱因斯坦求和约定，RAFT里的相关性体积，输出(B,H,W,W)
        corr = corr.reshape(B, H, W1, 1, W2).contiguous() # 改变张量形状(B,H,W,1,W)
        return corr #返回1/4原尺寸全对相关性体积
